# 用户个人的化身 根据提示词来扮演用户本人 模仿本人的思维、性格、兴趣爱好
# langchain

from langchain.agents import AgentExecutor
from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_openai import ChatOpenAI
from langchain.agents.format_scratchpad import format_to_openai_functions
from langchain.agents.output_parsers import OpenAIFunctionsAgentOutputParser
from langchain.tools import Tool, format_tool_to_openai_function
from langchain.memory import ConversationBufferMemory
from langchain.schema import HumanMessage, AIMessage
from typing import Dict, Any, List, Optional
from . import ai_base
from sqlalchemy.orm import Session
from .. import models, schemas
from ..crud import chat as chat_crud
from fastapi import HTTPException
import uuid

class AvatarAgent:
    def __init__(self, profile: Dict[str, str], db: Session = None, avatar_id: str = None, user_id: str = None):
        """
        初始化化身代理
        
        Args:
            profile: 包含以下字段的字典:
                - username: 用户名
                - personality: 性格特征
                - interests: 兴趣爱好
                - values: 价值观
                - speaking_style: 说话风格
                - sample_dialogue: 示例对话
            db: 数据库会话
            avatar_id: 化身ID
            user_id: 用户ID
        """
        self.profile = profile
        self.db = db
        self.avatar_id = avatar_id
        self.user_id = user_id
        self.llm = ai_base.llm
        self.tools = []
        
        # 创建或获取聊天会话
        if db and avatar_id and user_id:
            self.session = self._get_or_create_session()
            # 从数据库加载历史消息到memory
            self.memory = self._load_chat_history()
        else:
            self.memory = ConversationBufferMemory(
                return_messages=True,
                memory_key='chat_history'
            )
        
        # 构建提示模板
        self.prompt = ChatPromptTemplate.from_messages([
            ("system", f"""你是一个人,名字叫{profile['username']}。
            
            你的性格特征: {profile['personality']}
            你的兴趣爱好: {profile['interests']}
            你的价值观: {profile['values']}
            你的说话风格: {profile['speaking_style']}
            
            请始终保持这些特征，与用户对话时要符合这个人设，回答要口语化，每次回答要简短，不要超过20个字，不必在意语法正确性，你是一个网友。
            """),
            MessagesPlaceholder(variable_name="chat_history"),
            ("human", "{input}"),
            MessagesPlaceholder(variable_name="agent_scratchpad"),
        ])
        
        # 创建 agent
        llm_with_tools = self.llm.bind(
            functions=[format_tool_to_openai_function(t) for t in self.tools]
        )
        
        self.agent = (
            {
                "input": lambda x: x["input"],
                "chat_history": lambda x: self.memory.load_memory_variables({})["chat_history"],
                "agent_scratchpad": lambda x: format_to_openai_functions(x["intermediate_steps"])
            }
            | self.prompt
            | llm_with_tools
            | OpenAIFunctionsAgentOutputParser()
        )
        
        self.agent_executor = AgentExecutor(
            agent=self.agent,
            tools=self.tools,
            memory=self.memory,
            max_iterations=1,
            early_stopping_method="force",
            verbose=True,
            handle_parsing_errors=True,
            return_intermediate_steps=False
        )
    
    def _get_or_create_session(self) -> models.AvatarChatSession:
        """获取或创建聊天会话"""
        session = self.db.query(models.AvatarChatSession).filter(
            models.AvatarChatSession.user_id == self.user_id,
            models.AvatarChatSession.avatar_id == self.avatar_id
        ).first()
        
        if not session:
            session = models.AvatarChatSession(
                id=f"s_{uuid.uuid4().hex}",
                user_id=self.user_id,
                avatar_id=self.avatar_id
            )
            self.db.add(session)
            self.db.commit()
            self.db.refresh(session)
        
        return session
    
    def _load_chat_history(self) -> ConversationBufferMemory:
        """从数据库加载聊天历史"""
        memory = ConversationBufferMemory(
            return_messages=True,
            memory_key='chat_history'
        )
        
        # 获取当前会话最近的10条消息
        messages = self.db.query(models.AvatarChatMessage).filter(
            models.AvatarChatMessage.session_id == self.session.id  # 根据session_id过滤
        ).order_by(models.AvatarChatMessage.created_at.desc()).limit(10).all()
        
        # 反转消息顺序，从旧到新
        messages.reverse()
        
        # 将消息添加到memory
        for msg in messages:
            if msg.is_from_user:
                memory.chat_memory.add_message(HumanMessage(content=msg.content))
            else:
                memory.chat_memory.add_message(AIMessage(content=msg.content))
        
        return memory
    
    def _save_message(self, content: str, is_from_user: bool):
        """保存消息到数据库"""
        if not (self.db and self.avatar_id and self.user_id and self.session):
            return
            
        message = models.AvatarChatMessage(
            content=content,
            user_id=self.user_id,
            avatar_id=self.avatar_id,
            is_from_user=is_from_user,
            session_id=self.session.id
        )
        self.db.add(message)
        self.db.commit()
        self.db.refresh(message)
        return message
    
    def chat(self, user_input: str) -> str:
        """处理用户输入并返回化身回复"""
        try:
            # 保存用户消息
            self._save_message(user_input, is_from_user=True)
            
            # 获取AI回复
            response = self.agent_executor.invoke({
                "input": user_input
            })
            response_text = response["output"] if isinstance(response, dict) else str(response)
            
            # 保存AI回复
            self._save_message(response_text, is_from_user=False)
            
            return response_text
        except Exception as e:
            print(f"Avatar chat error: {str(e)}")
            return f"对话出现错误: {str(e)}"

    @classmethod
    def from_db(cls, db: Session, avatar_id: str, user_id: str, session_id: Optional[str] = None) -> "AvatarAgent":
        """从数据库创建AvatarAgent实例"""
        avatar = db.query(models.Avatar).filter(models.Avatar.id == avatar_id).first()
        if not avatar:
            raise HTTPException(status_code=404, detail="化身不存在")
        
        # 如果没有指定session_id，获取最新的会话
        if not session_id:
            session = chat_crud.get_latest_avatar_chat_session(db, user_id, avatar_id)
            if not session:
                session = chat_crud.create_avatar_chat_session(db, user_id, avatar_id)
        else:
            session = db.query(models.AvatarChatSession).get(session_id)
            if not session:
                raise HTTPException(status_code=404, detail="会话不存在")
        
        user = db.query(models.User).filter(models.User.id == avatar.user_id).first()
        if not user:
            raise HTTPException(status_code=404, detail="化身所属用户不存在")
        
        prompt_settings = dict(avatar.prompt_settings)
        prompt_settings["username"] = user.username
        
        required_fields = ["personality", "interests", "values", "speaking_style", "sample_dialogue"]
        missing_fields = [field for field in required_fields if field not in prompt_settings]
        
        if missing_fields:
            raise HTTPException(
                status_code=400, 
                detail=f"化身配置缺少必需字段: {', '.join(missing_fields)}"
            )
        
        try:
            instance = cls(prompt_settings, db, avatar_id, user_id)
            instance.session = session
            return instance
        except Exception as e:
            print(f"创建Agent失败，prompt_settings: {prompt_settings}")
            raise HTTPException(status_code=500, detail=f"创建化身Agent失败: {str(e)}")